"""
Core feature flag management system for CodeMCP.

Provides comprehensive feature flag management with runtime evaluation,
dependency resolution, and integration with existing infrastructure.
"""

import asyncio
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Set, Union
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

from ...core.config import Config
from ...core.error_handler import CodeMCPError, log_info, log_warning, log_debug, handle_error
from ...storage.migrations.migration_manager import MigrationManager


class FeatureState(Enum):
    """Feature flag states."""
    DISABLED = "disabled"
    ENABLED = "enabled"
    CONDITIONAL = "conditional"
    ROLLOUT = "rollout"
    DEPRECATED = "deprecated"


class TargetingRule(Enum):
    """Targeting rule types."""
    USER_ID = "user_id"
    USER_GROUP = "user_group"
    ENVIRONMENT = "environment"
    PERCENTAGE = "percentage"
    TIME_WINDOW = "time_window"
    CUSTOM = "custom"


@dataclass
class FeatureFlag:
    """Feature flag definition."""
    name: str
    description: str
    state: FeatureState = FeatureState.DISABLED
    environment: str = "development"
    rollout_percentage: float = 0.0
    targeting_rules: Dict[str, Any] = field(default_factory=dict)
    dependencies: Set[str] = field(default_factory=set)
    conflicts: Set[str] = field(default_factory=set)
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.utcnow)
    updated_at: datetime = field(default_factory=datetime.utcnow)
    version: int = 1
    tags: Set[str] = field(default_factory=set)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert feature flag to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "state": self.state.value,
            "environment": self.environment,
            "rollout_percentage": self.rollout_percentage,
            "targeting_rules": self.targeting_rules,
            "dependencies": list(self.dependencies),
            "conflicts": list(self.conflicts),
            "metadata": self.metadata,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "version": self.version,
            "tags": list(self.tags)
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "FeatureFlag":
        """Create feature flag from dictionary."""
        return cls(
            name=data["name"],
            description=data["description"],
            state=FeatureState(data.get("state", "disabled")),
            environment=data.get("environment", "development"),
            rollout_percentage=data.get("rollout_percentage", 0.0),
            targeting_rules=data.get("targeting_rules", {}),
            dependencies=set(data.get("dependencies", [])),
            conflicts=set(data.get("conflicts", [])),
            metadata=data.get("metadata", {}),
            created_at=datetime.fromisoformat(data.get("created_at", datetime.utcnow().isoformat())),
            updated_at=datetime.fromisoformat(data.get("updated_at", datetime.utcnow().isoformat())),
            version=data.get("version", 1),
            tags=set(data.get("tags", []))
        )


@dataclass
class EvaluationContext:
    """Context for feature flag evaluation."""
    user_id: Optional[str] = None
    user_groups: Set[str] = field(default_factory=set)
    environment: str = "development"
    custom_attributes: Dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=datetime.utcnow)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert evaluation context to dictionary."""
        return {
            "user_id": self.user_id,
            "user_groups": list(self.user_groups),
            "environment": self.environment,
            "custom_attributes": self.custom_attributes,
            "timestamp": self.timestamp.isoformat()
        }


class FeatureFlagError(CodeMCPError):
    """Feature flag specific errors."""
    
    def __init__(self, message: str, flag_name: str = None, **kwargs):
        super().__init__(message, "FEATURE_FLAG_ERROR", **kwargs)
        if flag_name:
            self.details["flag_name"] = flag_name


class FeatureFlagManager:
    """
    Core feature flag manager for CodeMCP.
    
    Provides comprehensive feature flag management with runtime evaluation,
    dependency resolution, and integration with migration infrastructure.
    """
    
    def __init__(self, config: Config = None, migration_manager: MigrationManager = None):
        self.config = config or Config()
        self.migration_manager = migration_manager
        
        # Feature flag registry
        self.flags: Dict[str, FeatureFlag] = {}
        self.flag_cache: Dict[str, Dict[str, Any]] = {}
        
        # Performance tracking
        self.evaluation_stats = {
            "total_evaluations": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "avg_evaluation_time": 0.0,
            "evaluation_times": []
        }
        
        # Configuration
        self.cache_ttl_seconds = 300  # 5 minutes
        self.max_cache_size = 10000
        self.enable_performance_tracking = True
        self.evaluation_timeout_seconds = 1.0
        
        log_info("Feature flag manager initialized")
    
    async def initialize(self):
        """Initialize the feature flag manager."""
        try:
            # Initialize migration manager if provided
            if self.migration_manager and not self.migration_manager.script_directory:
                await self.migration_manager.initialize()
            
            # Create feature flag schema if needed
            await self._ensure_schema_exists()
            
            # Load existing flags
            await self._load_flags_from_storage()
            
            log_info("Feature flag manager initialized successfully")
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "feature_flag_manager_init"})
            raise FeatureFlagError(f"Failed to initialize feature flag manager: {error_response['message']}")
    
    async def _ensure_schema_exists(self):
        """Ensure feature flag database schema exists."""
        if not self.migration_manager:
            log_warning("No migration manager provided, skipping schema creation")
            return
        
        try:
            # Check if feature flag tables exist by attempting a migration
            schema_migration = """
-- Create feature flags table
CREATE TABLE IF NOT EXISTS feature_flags (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name VARCHAR(255) NOT NULL UNIQUE,
    description TEXT,
    state VARCHAR(50) NOT NULL DEFAULT 'disabled',
    environment VARCHAR(100) NOT NULL DEFAULT 'development',
    rollout_percentage REAL DEFAULT 0.0,
    targeting_rules TEXT, -- JSON
    dependencies TEXT, -- JSON array
    conflicts TEXT, -- JSON array
    metadata TEXT, -- JSON
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    version INTEGER DEFAULT 1,
    tags TEXT -- JSON array
);

-- Create feature flag evaluations table for caching
CREATE TABLE IF NOT EXISTS feature_flag_evaluations (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    flag_name VARCHAR(255) NOT NULL,
    context_hash VARCHAR(64) NOT NULL,
    result BOOLEAN NOT NULL,
    evaluation_time REAL NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    expires_at TIMESTAMP NOT NULL,
    UNIQUE(flag_name, context_hash)
);

-- Create indexes for performance
CREATE INDEX IF NOT EXISTS idx_feature_flags_name ON feature_flags(name);
CREATE INDEX IF NOT EXISTS idx_feature_flags_environment ON feature_flags(environment);
CREATE INDEX IF NOT EXISTS idx_feature_flags_state ON feature_flags(state);
CREATE INDEX IF NOT EXISTS idx_feature_flag_evaluations_expires ON feature_flag_evaluations(expires_at);
"""
            
            # Create migration if needed
            migration_revision = await self.migration_manager.create_migration(
                "Add feature flag schema",
                autogenerate=False
            )
            
            if migration_revision:
                log_info(f"Created feature flag schema migration: {migration_revision}")
            
        except Exception as e:
            log_warning(f"Failed to create feature flag schema: {e}")
    
    async def _load_flags_from_storage(self):
        """Load existing feature flags from storage."""
        # For now, this is a placeholder. In a real implementation,
        # this would load from the database or configuration files.
        try:
            # Load default feature flags for CodeMCP GSI Wiki Generator
            default_flags = {
                "enhanced_code_analysis": FeatureFlag(
                    name="enhanced_code_analysis",
                    description="Enhanced code analysis with advanced pattern recognition",
                    state=FeatureState.CONDITIONAL,
                    targeting_rules={
                        "environments": ["development", "staging"],
                        "user_groups": ["beta_users", "developers"]
                    },
                    tags={"core", "analysis"}
                ),
                "gsi_drp4dvar_integration": FeatureFlag(
                    name="gsi_drp4dvar_integration",
                    description="DRP-4DVar integration for GSI analysis",
                    state=FeatureState.ROLLOUT,
                    rollout_percentage=50.0,
                    targeting_rules={
                        "environments": ["development", "staging"]
                    },
                    tags={"gsi", "integration"}
                ),
                "real_time_wiki_updates": FeatureFlag(
                    name="real_time_wiki_updates",
                    description="Real-time wiki updates and live collaboration",
                    state=FeatureState.DISABLED,
                    dependencies={"enhanced_code_analysis"},
                    tags={"wiki", "real-time"}
                ),
                "advanced_caching": FeatureFlag(
                    name="advanced_caching",
                    description="Advanced caching system for improved performance",
                    state=FeatureState.ENABLED,
                    targeting_rules={
                        "environments": ["production", "staging"]
                    },
                    tags={"performance", "caching"}
                ),
                "ml_code_suggestions": FeatureFlag(
                    name="ml_code_suggestions",
                    description="Machine learning-powered code suggestions",
                    state=FeatureState.DISABLED,
                    conflicts={"enhanced_code_analysis"},
                    tags={"ai", "experimental"}
                )
            }
            
            self.flags.update(default_flags)
            log_info(f"Loaded {len(default_flags)} default feature flags")
            
        except Exception as e:
            log_warning(f"Failed to load flags from storage: {e}")
    
    async def register_flag(self, flag: FeatureFlag) -> bool:
        """Register a new feature flag."""
        try:
            if flag.name in self.flags:
                existing_flag = self.flags[flag.name]
                flag.version = existing_flag.version + 1
                flag.updated_at = datetime.utcnow()
                log_info(f"Updated existing feature flag: {flag.name} (v{flag.version})")
            else:
                log_info(f"Registered new feature flag: {flag.name}")
            
            self.flags[flag.name] = flag
            
            # Clear cache for this flag
            await self._clear_flag_cache(flag.name)
            
            return True
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "register_flag", "flag_name": flag.name})
            raise FeatureFlagError(f"Failed to register flag: {error_response['message']}", flag.name)
    
    async def unregister_flag(self, flag_name: str) -> bool:
        """Unregister a feature flag."""
        try:
            if flag_name not in self.flags:
                log_warning(f"Flag not found: {flag_name}")
                return False
            
            del self.flags[flag_name]
            await self._clear_flag_cache(flag_name)
            
            log_info(f"Unregistered feature flag: {flag_name}")
            return True
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "unregister_flag", "flag_name": flag_name})
            raise FeatureFlagError(f"Failed to unregister flag: {error_response['message']}", flag_name)
    
    async def evaluate_flag(self, flag_name: str, context: EvaluationContext) -> bool:
        """Evaluate a feature flag for the given context."""
        start_time = time.time()
        
        try:
            # Check if flag exists
            if flag_name not in self.flags:
                log_debug(f"Flag not found: {flag_name}, returning False")
                return False
            
            flag = self.flags[flag_name]
            
            # Check cache first
            cache_key = self._generate_cache_key(flag_name, context)
            cached_result = await self._get_cached_evaluation(cache_key)
            if cached_result is not None:
                self.evaluation_stats["cache_hits"] += 1
                return cached_result
            
            self.evaluation_stats["cache_misses"] += 1
            
            # Perform evaluation
            result = await self._evaluate_flag_logic(flag, context)
            
            # Cache the result
            await self._cache_evaluation(cache_key, result, time.time() - start_time)
            
            return result
            
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "evaluate_flag",
                "flag_name": flag_name,
                "context": context.to_dict()
            })
            log_warning(f"Flag evaluation failed: {error_response['message']}")
            return False
            
        finally:
            # Update performance stats
            if self.enable_performance_tracking:
                evaluation_time = time.time() - start_time
                self._update_evaluation_stats(evaluation_time)
    
    async def _evaluate_flag_logic(self, flag: FeatureFlag, context: EvaluationContext) -> bool:
        """Core flag evaluation logic."""
        # Check basic state
        if flag.state == FeatureState.DISABLED:
            return False
        
        if flag.state == FeatureState.ENABLED:
            # Still need to check environment and targeting
            pass
        
        # Check environment match
        if context.environment != flag.environment and flag.environment != "all":
            return False
        
        # Check dependencies
        for dep_flag in flag.dependencies:
            dep_result = await self.evaluate_flag(dep_flag, context)
            if not dep_result:
                log_debug(f"Dependency {dep_flag} not satisfied for {flag.name}")
                return False
        
        # Check conflicts
        for conflict_flag in flag.conflicts:
            conflict_result = await self.evaluate_flag(conflict_flag, context)
            if conflict_result:
                log_debug(f"Conflict {conflict_flag} active for {flag.name}")
                return False
        
        # Evaluate targeting rules
        if not await self._evaluate_targeting_rules(flag, context):
            return False
        
        # Check rollout percentage
        if flag.rollout_percentage < 100.0:
            if not await self._evaluate_rollout(flag, context):
                return False
        
        return True
    
    async def _evaluate_targeting_rules(self, flag: FeatureFlag, context: EvaluationContext) -> bool:
        """Evaluate targeting rules for a flag."""
        if not flag.targeting_rules:
            return True
        
        try:
            # Environment targeting
            if "environments" in flag.targeting_rules:
                allowed_envs = flag.targeting_rules["environments"]
                if context.environment not in allowed_envs:
                    return False
            
            # User ID targeting
            if "user_ids" in flag.targeting_rules and context.user_id:
                allowed_users = flag.targeting_rules["user_ids"]
                if context.user_id not in allowed_users:
                    return False
            
            # User group targeting
            if "user_groups" in flag.targeting_rules and context.user_groups:
                required_groups = set(flag.targeting_rules["user_groups"])
                if not required_groups.intersection(context.user_groups):
                    return False
            
            # Time window targeting
            if "time_window" in flag.targeting_rules:
                time_window = flag.targeting_rules["time_window"]
                if not await self._evaluate_time_window(time_window, context):
                    return False
            
            # Custom attribute targeting
            if "custom_attributes" in flag.targeting_rules:
                custom_rules = flag.targeting_rules["custom_attributes"]
                if not await self._evaluate_custom_attributes(custom_rules, context):
                    return False
            
            return True
            
        except Exception as e:
            log_warning(f"Failed to evaluate targeting rules for {flag.name}: {e}")
            return False
    
    async def _evaluate_rollout(self, flag: FeatureFlag, context: EvaluationContext) -> bool:
        """Evaluate rollout percentage for a flag."""
        if flag.rollout_percentage >= 100.0:
            return True
        
        if flag.rollout_percentage <= 0.0:
            return False
        
        # Use consistent hashing based on flag name and user context
        hash_input = f"{flag.name}:{context.user_id or 'anonymous'}:{context.environment}"
        hash_value = hash(hash_input) % 100
        
        return hash_value < flag.rollout_percentage
    
    async def _evaluate_time_window(self, time_window: Dict[str, Any], context: EvaluationContext) -> bool:
        """Evaluate time window targeting."""
        try:
            now = context.timestamp
            
            if "start" in time_window:
                start_time = datetime.fromisoformat(time_window["start"])
                if now < start_time:
                    return False
            
            if "end" in time_window:
                end_time = datetime.fromisoformat(time_window["end"])
                if now > end_time:
                    return False
            
            return True
            
        except Exception as e:
            log_warning(f"Failed to evaluate time window: {e}")
            return False
    
    async def _evaluate_custom_attributes(self, custom_rules: Dict[str, Any], context: EvaluationContext) -> bool:
        """Evaluate custom attribute targeting."""
        try:
            for attr_name, expected_value in custom_rules.items():
                if attr_name not in context.custom_attributes:
                    return False
                
                actual_value = context.custom_attributes[attr_name]
                
                # Support different comparison types
                if isinstance(expected_value, dict):
                    if "equals" in expected_value:
                        if actual_value != expected_value["equals"]:
                            return False
                    elif "contains" in expected_value:
                        if expected_value["contains"] not in str(actual_value):
                            return False
                    elif "greater_than" in expected_value:
                        if float(actual_value) <= float(expected_value["greater_than"]):
                            return False
                else:
                    # Simple equality check
                    if actual_value != expected_value:
                        return False
            
            return True
            
        except Exception as e:
            log_warning(f"Failed to evaluate custom attributes: {e}")
            return False
    
    def _generate_cache_key(self, flag_name: str, context: EvaluationContext) -> str:
        """Generate cache key for evaluation."""
        import hashlib
        
        context_str = f"{flag_name}:{context.user_id}:{context.environment}:{sorted(context.user_groups)}:{context.custom_attributes}"
        return hashlib.sha256(context_str.encode()).hexdigest()[:16]
    
    async def _get_cached_evaluation(self, cache_key: str) -> Optional[bool]:
        """Get cached evaluation result."""
        try:
            if cache_key in self.flag_cache:
                cache_entry = self.flag_cache[cache_key]
                if time.time() < cache_entry["expires_at"]:
                    return cache_entry["result"]
                else:
                    # Remove expired entry
                    del self.flag_cache[cache_key]
            
            return None
            
        except Exception as e:
            log_debug(f"Failed to get cached evaluation: {e}")
            return None
    
    async def _cache_evaluation(self, cache_key: str, result: bool, evaluation_time: float):
        """Cache evaluation result."""
        try:
            # Clean cache if it's too large
            if len(self.flag_cache) >= self.max_cache_size:
                await self._clean_cache()
            
            expires_at = time.time() + self.cache_ttl_seconds
            self.flag_cache[cache_key] = {
                "result": result,
                "evaluation_time": evaluation_time,
                "expires_at": expires_at
            }
            
        except Exception as e:
            log_debug(f"Failed to cache evaluation: {e}")
    
    async def _clean_cache(self):
        """Clean expired cache entries."""
        try:
            current_time = time.time()
            expired_keys = [
                key for key, entry in self.flag_cache.items()
                if current_time >= entry["expires_at"]
            ]
            
            for key in expired_keys:
                del self.flag_cache[key]
            
            # If still too large, remove oldest entries
            if len(self.flag_cache) >= self.max_cache_size:
                sorted_entries = sorted(
                    self.flag_cache.items(),
                    key=lambda x: x[1]["expires_at"]
                )
                
                # Remove oldest 25%
                num_to_remove = len(sorted_entries) // 4
                for key, _ in sorted_entries[:num_to_remove]:
                    del self.flag_cache[key]
            
        except Exception as e:
            log_warning(f"Failed to clean cache: {e}")
    
    async def _clear_flag_cache(self, flag_name: str):
        """Clear cache entries for a specific flag."""
        try:
            keys_to_remove = [
                key for key in self.flag_cache.keys()
                if key.startswith(f"{flag_name}:")
            ]
            
            for key in keys_to_remove:
                del self.flag_cache[key]
            
        except Exception as e:
            log_debug(f"Failed to clear flag cache: {e}")
    
    def _update_evaluation_stats(self, evaluation_time: float):
        """Update evaluation performance statistics."""
        try:
            self.evaluation_stats["total_evaluations"] += 1
            self.evaluation_stats["evaluation_times"].append(evaluation_time)
            
            # Keep only last 1000 evaluation times for moving average
            if len(self.evaluation_stats["evaluation_times"]) > 1000:
                self.evaluation_stats["evaluation_times"] = self.evaluation_stats["evaluation_times"][-1000:]
            
            # Update average
            self.evaluation_stats["avg_evaluation_time"] = (
                sum(self.evaluation_stats["evaluation_times"]) / 
                len(self.evaluation_stats["evaluation_times"])
            )
            
        except Exception as e:
            log_debug(f"Failed to update evaluation stats: {e}")
    
    async def get_flag(self, flag_name: str) -> Optional[FeatureFlag]:
        """Get a feature flag by name."""
        return self.flags.get(flag_name)
    
    async def list_flags(self, environment: str = None, state: FeatureState = None) -> List[FeatureFlag]:
        """List feature flags with optional filtering."""
        flags = list(self.flags.values())
        
        if environment:
            flags = [f for f in flags if f.environment == environment or f.environment == "all"]
        
        if state:
            flags = [f for f in flags if f.state == state]
        
        return flags
    
    async def get_evaluation_stats(self) -> Dict[str, Any]:
        """Get evaluation performance statistics."""
        stats = self.evaluation_stats.copy()
        stats["cache_size"] = len(self.flag_cache)
        stats["cache_hit_rate"] = (
            self.evaluation_stats["cache_hits"] / 
            max(1, self.evaluation_stats["cache_hits"] + self.evaluation_stats["cache_misses"])
        )
        
        return stats
    
    async def health_check(self) -> Dict[str, Any]:
        """Perform health check of the feature flag system."""
        try:
            # Basic health metrics
            health = {
                "status": "healthy",
                "flags_registered": len(self.flags),
                "cache_size": len(self.flag_cache),
                "total_evaluations": self.evaluation_stats["total_evaluations"],
                "avg_evaluation_time": self.evaluation_stats["avg_evaluation_time"],
                "cache_hit_rate": (
                    self.evaluation_stats["cache_hits"] / 
                    max(1, self.evaluation_stats["cache_hits"] + self.evaluation_stats["cache_misses"])
                )
            }
            
            # Check if any critical flags are in error state
            critical_flags = [f for f in self.flags.values() if "critical" in f.tags]
            health["critical_flags"] = len(critical_flags)
            
            return health
            
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }